# Author : ZZH
# Date : 2025/5/20
import yaml
from pathlib import Path
import logging
import random
from datetime import datetime


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s -%(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def verify_dataset_config(yaml_path, datailed=False, sample_radio=0.1, min_sample=10):
    """
    验证YOLO数据集配置，检查data.yaml和对应图像，标签文件
    :param yaml_path: data.yaml路径
    :param datailed: 是否进行全量验证
    :param sample_radio: 抽样比例
    :param min_sample: 最小抽样数量
    :return: 验证通过返回True，否则False
    """
    yaml_path = Path(yaml_path).resolve()
    logger.info(f"验证data.yaml路径为{yaml_path}")
    if not yaml_path.exists():
        logger.error(f"data.yaml文件不存在:{yaml_path}")
        return False
    with open(yaml_path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)

    class_names = config.get('names', [])
    nc = config.get('nc', 0)

    if len(class_names) != nc:
        logger.error(f"class_names和nc长度不匹配:{len(class_names)}不等于{nc}")
        return False
    logger.error(f"class_names和nc长度匹配，类比数量:{nc}, 类别为；{class_names}")

    # 验证数据集（train,val,test)
    # test可选，存在就验证，否则验证train和val
    splits = ['train', 'val', 'test'] if 'test' in config else ['train', 'val']
    for split in splits:
        split_path = (yaml_path.parent / config[split]).resolve()
        logger.info(f"验证{split}路径为:{split_path}")
        if not split_path.exists():
            logger.error(f"图像目录{split_path}路径不存在")
            return False
        img_paths = list(split_path.glob('*.[jJ][pP][gG]')) + \
                    list(split_path.glob('*.[pP][nN][gG]')) + \
                    list(split_path.glob('*.[jJ][pP][eE][gG]')) + \
                    list(split_path.glob('*.[bB][mM][pP]'))
        if not img_paths:
            logger.error(f"图像目录{split_path}下没有图像文件")
            return False
        logger.info(f"{split}路径下图像数量为:{len(img_paths)}")

        # 动态抽样大小
        sample_size = max(min_sample, int(len(img_paths) * sample_radio))

        if datailed:
            logger.info(f"{split}进行全量验证,抽样数量为:{sample_size}")
            sample_paths = img_paths
        else:
            logger.info(f"{split}进行抽样验证，抽样数量为:{sample_size}")
            sample_paths = random.sample(img_paths, min(sample_size, len(img_paths)))
        for img_path in sample_paths:
            img_path_resolve = img_path.resolve()
            logger.debug(f"验证图像文件: {img_path_resolve}")
            if not img_path_resolve.exists():
                logger.error(f"图像文件{img_path_resolve} 不存在")
                return False
            label_path = split_path.parent / split / (img_path.stem + '.txt')
            logger.debug(f"验证标签文件:{label_path}")

            if not label_path.exists():
                logger.debug(f"标签文件{label_path}不存在(正常,空标签)")
                continue
            # 验证Yolo格式的标签文件内容
            with open(label_path, 'r', encoding='utf-8') as f:
                lines = f.read().splitlines()
            if not lines:
                logger.debug(f"标签文件{label_path}为空.(正常的)")
                continue
            for line in lines:
                parts = line.split()
                if len(parts) != 5:
                    logger.error(f"{label_path} 内容格式错误: {line}")
                    return False
                try:
                    class_id = int(parts[0])
                    if class_id < 0 or class_id >= nc:
                        logger.error(f"标签文件{label_path} 内容错误:类别id{class_id} 超出nc {nc}范围")
                        return False
                    coords = [float(x) for x in parts[1:]]
                    if not all(0 <= x <= 1 for x in coords):
                        logger.error(f"标签文件{label_path} 内容错误:坐标{coords}超出范围 [0, 1]: {line}")
                        return False
                except ValueError:
                    logger.error(f"标签文件{label_path}包含无效值: {line}")
                    return False

    logger.info(f"数据集验证通过")
    return True


if __name__ == '__main__':
    base_path = Path(__file__).resolve().parent.parent      # 项目根目录
    logger.info(f"项目根目录:{base_path}")
    yaml_path = base_path / "config" / "data.yaml"

    # 抽样验证
    if verify_dataset_config(yaml_path, datailed=False, sample_radio=0.1, min_sample=10):
        logger.info(f"数据集验证通过")
    else:
        logger.error(f"数据集验证失败")

    # 全局验证
    if verify_dataset_config(yaml_path, datailed=True):
        logger.info(f"数据集验证通过")
    else:
        logger.error(f"数据集验证失败")
